import os.path
import Queue
import threading
import multiprocessing
import re

from doas import io, spectrum_maths
from doas.watcher import create_dir_watcher
from std_ops.os_ import find_files
from std_ops.iter_ import list_splitter


#TODO - spectrum metadata loader (remove the getSpectrometer name etc from the 
#spectrum directory class). Could also have a realtime metadata loader?
#TODO - add load_meta_data method to spectrum plugins.

__seq_num_regex = re.compile(".*?([0-9]{1,})(?:\..*){0,1}$")

def spec_obj_name_cmp(s1, s2):
    return spec_name_cmp(s1.filename, s2.filename)

def spec_name_cmp(s1, s2):
    """
    Comparator for sorting spectra into the order in which they were
    captured. We expect that the filenames will consist of some prefix,
    followed by a sequence number and then a filename extension. This
    comparator sorts the spectra by their sequence number. No assumptions
    are made about zero padding in the sequence number.
    
    If no sequence number can be extracted from one of the filenames then 
    they are compared by value instead. 
    """
    seq_num1 = __seq_num_regex.match(s1)
    if seq_num1 is None:
        return cmp(s1, s2)
    
    seq_num2 = __seq_num_regex.match(s2)
    if seq_num2 is None:
        return cmp(s1, s2)
    
    seq_num1 = int(seq_num1.groups()[0])
    seq_num2 = int(seq_num2.groups()[0])
    
    return cmp(seq_num1, seq_num2)


def _async_load(arg_tuple):
    loader_obj, spectrum_file = arg_tuple
    try:
        s = loader_obj.load(spectrum_file)
        return s
    except io.UnableToLoad:
        return None


def _async_load_with_dark_calc(arg_tuple):
    loader_obj, spectrum_file = arg_tuple
    try:
        s = loader_obj.load(spectrum_file)
        s.is_dark()
        return s
    except io.UnableToLoad:
        return None


class SpecDirIter:
    def __init__(self, directory, realtime=False, skip_existing=False, recursive=False, chunk_size=100):
        
        if not os.path.isdir(directory):
            raise ValueError("Cannot access "+directory+". No such directory.")
        
        if skip_existing and not realtime:
            raise RuntimeError("If skip_existing is set to True then realtime must also be set to True (otherwise no spectra will be returned).")
        
        self.__dir = directory
        self._spectra_obj_q = Queue.Queue(maxsize=3000)
        self.__realtime_filename_q = Queue.Queue()
        self._stay_alive = True
        self.__chunk_size = chunk_size
        self.__existing_loader_thread = None
        self._realtime_loader_thread = None
        self._finished_loading_existing_lock = threading.Lock()
        self._finished_loading_existing_lock.acquire()
        self.__existing_spectra_found = {}
        self.__recursive = recursive
        self.__num_of_existing_to_load = None
        self.__existing_files_metadata_lock = threading.Lock()
        self.__existing_files_metadata_lock.acquire()
        
        #don't want to overwrite potential subclass async load function value if
        #this constructor is called after it has been set in the subclass
        if not hasattr(self, '_async_load_function'):
            self._async_load_function = _async_load
        
        #note that we start the realtime loader thread BEFORE the existing loader thread
        #this is to prevent spectra that are created during the find_files() call being 
        #skipped.
        if realtime:
            self._realtime_loader_thread = threading.Thread(target=self.__load_realtime)
            self._realtime_loader_thread.start()
              
        if not skip_existing:
            self.__existing_loader_thread = threading.Thread(target=self.__load_existing)
            self.__existing_loader_thread.start()
        else:
            self._finished_loading_existing_lock.release()
            self.__existing_files_metadata_lock.release()

           
    def __iter__(self):
        """
        Method required by iterator protocol. Allows iterator to be used in 
        for loops.
        """
        return self

    def __next__(self):
        #needed for Py3k compatibility
        return self.next()
        
        
    def next(self):

        s = self._spectra_obj_q.get(block=True)
        
        if s is None or not self._stay_alive:
            
            raise StopIteration
        
        return s

    
    def close(self):
        self._stay_alive = False
        
        #the loader threads may be blocking waiting to put something into the queue
        while True:
            try:
                self._spectra_obj_q.get(block=False)
            except Queue.Empty:
                break
            
        #the next() method may be blocking waiting to get something from the queue
        try:
            self._spectra_obj_q.put(None, block=False)
        except Queue.Full:
            pass
        
        #the realtime loader thread may be blocking waiting to get a new filename to load
        try:
            self.__realtime_filename_q.put(None)
        except Queue.Full:
            pass
                  
        if self.__existing_loader_thread is not None:
            self.__existing_loader_thread.join()
            
        if self._realtime_loader_thread is not None:
            self._realtime_loader_thread.join()
               
                
    def get_num_existing(self):
        if self.__existing_loader_thread is None:
            raise RuntimeError("Cannot call get_num_existing if skip_existing flag is set.")
               
        self.__existing_files_metadata_lock.acquire()
        self.__existing_files_metadata_lock.release()
        return self.__num_of_existing_to_load
    
    def get_spectrometer_name(self):
        if self.__existing_loader_thread is None:
            raise RuntimeError("Cannot call get_spectrometer_name if skip_existing flag is set.")
   
        self.__existing_files_metadata_lock.acquire()
        self.__existing_files_metadata_lock.release()
        return self.__spectrometer_name
    
       
    def __load_existing(self):
        try:
            loader = io.SpectrumIO()
            
            #find_files is slow so use listdir if we don't need to be recursive
            if self.__recursive:
                existing_files = find_files(self.__dir, recursive=True)
            else:
                existing_files = [os.path.join(self.__dir, n) for n in os.listdir(self.__dir)]
            
            self.__num_of_existing_to_load = len(existing_files)
                
            #sort the filenames by sequence number
            existing_files.sort(cmp=spec_name_cmp)
            
            for spectrum_file in existing_files:
                self.__existing_spectra_found[spectrum_file] = None
            
            try:
                n_cpus = multiprocessing.cpu_count()
            except NotImplemented:
                n_cpus = 1
            
            #load the first spectrum that we can to find out which spectrometer we are dealing with
            self.__spectrometer_name = None
            for f in existing_files:
                try:
                    s = loader.load(f)
                    self.__spectrometer_name = s.spectrometer_id
                    break
                except io.UnableToLoad:
                    continue
            if self.__spectrometer_name == None:
                #if we can't find the spectrometer name, then it is probably
                #because there are no valid spectra in the directory yet
                return
            self.__existing_files_metadata_lock.release()
            
            if not self._stay_alive:
                return 
                       
            if self.__num_of_existing_to_load < 300 or n_cpus == 1:
                #don't bother trying to parallel process the spectra - there aren't enough
                #to make it worth while
                #TODO - progress feedback
                for spectrum_file in existing_files:
                    
                    if not self._stay_alive:
                        break
                                    
                    try:
                        s = loader.load(spectrum_file)
                    except io.UnableToLoad:
                        continue
                    
                    self._spectra_obj_q.put(s)
            else:
                #use multiple processing to speed up the loading
                p = multiprocessing.Pool(n_cpus)
                
                for progress,filenames in list_splitter(existing_files, self.__chunk_size):
                    result = p.map_async(self._async_load_function, zip([loader]*len(filenames),filenames))
                    
                    spectra = [s for s in result.get() if s is not None]
                    if not self._stay_alive:
                        p.close()
                        return 
                    
                    spectra.sort(cmp=spec_obj_name_cmp)
                    
                    for s in spectra:
                        self._spectra_obj_q.put(s) 
                                
        finally:
            self._finished_loading_existing_lock.release()
            
        if self._realtime_loader_thread is None:        
            #put None into the queue so that the iteration finishes when the q is emptied
            self._spectra_obj_q.put(None)
            
    
    
    def __load_realtime(self):
        
        put_into_q = lambda t,f,q: q.put(f)
        
        #TODO - implement recursive realtime loading 
        #combine watcher code with AvoScan watcher code       
        watcher = create_dir_watcher(self.__dir, put_into_q, self.__recursive, self.__realtime_filename_q)
        watcher.start()
        loader = io.SpectrumIO()
        
        #wait for the existing spectra to be loaded
        self._finished_loading_existing_lock.acquire()
        self._finished_loading_existing_lock.release()
        
        #skip any spectra that were picked up by the existing loader
        while self._stay_alive:
            spectrum_file = self.__realtime_filename_q.get(block=True)
            
            #note that None is not a key in the exisitng_spectra_found, so there is no
            #need to check for a None explicitly
            if not self.__existing_spectra_found.has_key(spectrum_file):
                break
        
        self.__existing_spectra_found = {} #done with this dict now - leave it to be GC'd
        
        while self._stay_alive:         
            if spectrum_file is None:
                break

            try:
                s = loader.load(spectrum_file)
                print "Loaded file %s" %spectrum_file
                self._spectra_obj_q.put(s, block=False)
            except io.UnableToLoad:
                pass
            except Queue.Full:
                #if the output queue is full, then just skip the spectrum - otherwise it might block
                #forever
                print "Warning! Spectrum loader output queue is full - skipping spectrum \'"+spectrum_file+"\'"
                   
            spectrum_file = self.__realtime_filename_q.get(block=True)
        
        watcher.stop()



class DarkSubtractedSpectra(SpecDirIter):
       
    def __init__(self, directory, realtime=False, skip_existing=False, recursive=False):
        self.__bkgd_spectrum = None
        self.__recieved_none = False
        self.__waiting_for_bkgd = []
        self._async_load_function = _async_load_with_dark_calc
        SpecDirIter.__init__(self, directory, realtime, skip_existing, recursive=recursive)
 
    
    def set_dark_spectra(self, spectra):
        """
        Manually set the dark spectra to be used for the subtraction.
        Spectra should be an iterable containing dark spectra objects.
        """
        
        #TODO - is it a problem that these darks might get interleaved with actual
        #spectra? It's not at the moment, but it might be in the future!
        for s in spectra:
            s.set_is_dark(True)
            print "put dark spectra %s into object queue"%s.filename
            self._spectra_obj_q.put(s)
        
        
        #TODO - is this actually threadsafe....too hard to think about at the moment
        if self.__recieved_none:
            self._spectra_obj_q.put(None)
        
        
       
    def next(self):
        if not self._stay_alive:
            raise StopIteration
        
        if self.__bkgd_spectrum is not None and len(self.__waiting_for_bkgd) > 0:
            return self.__waiting_for_bkgd.pop(0) - self.__bkgd_spectrum
        
        while True:
            s = self._spectra_obj_q.get(block=True)
            
            if not self._stay_alive:
                raise StopIteration
            
            if s is None: 
                if self.__bkgd_spectrum is not None:
                #only break out of this loop if we have already found a 
                #dark spectrum. Otherwise subsequent calls to set_dark_spectra
                #will be ignored
                    break
                else:
                    self.__recieved_none = True
                    continue
               
            if s.is_dark():
                print "found dark spectrum: ",s.filename
                if self.__bkgd_spectrum is None:
                    self.__bkgd_spectrum = s
                else:
                    self.__bkgd_spectrum = spectrum_maths.mean([self.__bkgd_spectrum, s])
            
            else:
                if self.__bkgd_spectrum is None:
                    self.__waiting_for_bkgd.append(s)
                elif len(self.__waiting_for_bkgd) > 0:
                    self.__waiting_for_bkgd.append(s)
                    return self.__waiting_for_bkgd.pop(0) - self.__bkgd_spectrum
                else:
                    return s - self.__bkgd_spectrum
                  
        if self.__bkgd_spectrum is not None and len(self.__waiting_for_bkgd) > 0:
            self._spectra_obj_q.put(None)
            return self.__waiting_for_bkgd.pop(0) - self.__bkgd_spectrum
        
        raise StopIteration
    
